#%%
import numpy as np, matplotlib.pyplot as plt, seaborn as sns, numpy as np, sys, os, tqdm, lmfit as lm, scipy.special as ss
import pandas as pd, scipy.stats as stats, scipy.optimize as op, scipy.interpolate as ip, pathlib
%matplotlib inline
%config InlineBackend.figure_format='retina'
plt.style.use('seaborn-whitegrid')
sns.set_context('paper')
sns.set_style('ticks')
cm = 1/2.54

#%%
# Load dietrich data
df_diet = pd.read_csv('other_shape_data/dietrich_settling_data.csv', header=1)
plt.loglog(df_diet.X, df_diet.Y, 'o')

x_data = np.log10(df_diet.X.values)
y_data = np.log10(df_diet.Y.values)

funco = lambda x, a, b, c, d, e: a + b*x + c*x**2 + d*x**3 + e*x**4
model = lm.Model(funco)
model.set_param_hint('a', value=1, vary=True)
model.set_param_hint('b', value=1, vary=True)
model.set_param_hint('c', value=1, vary=True)
model.set_param_hint('d', value=1, vary=True)
model.set_param_hint('e', value=1, vary=True)
params = model.make_params()

results = model.fit(y_data[np.isfinite(x_data)], params, x=x_data[np.isfinite(x_data)])
a = results.result.params['a'].value
b = results.result.params['b'].value
c = results.result.params['c'].value
d = results.result.params['d'].value
e = results.result.params['e'].value

x = np.logspace(-1,11,100)
plt.plot(x, 10**funco(np.log10(x), a, b, c, d, e), 'r-')

# func = lambda x: a + b*x - c*x**2 - d*x**3 + e*x**4
func = lambda x: -3.76715 + 1.92944*x - 0.09815*x**2 - 0.00575*x**3 + 0.00056*x**4
x = np.logspace(-1,11,100)
plt.plot(x, 10**func(np.log10(x)), 'k--')
# plt.xlim(2e5, 2e10)
# plt.ylim(100, 5e5)

xf = pd.DataFrame({'X_10': df_diet.X[df_diet.X > 2e5], 'Y_10': df_diet.Y[df_diet.X > 2e5], 'Yp_10': (df_diet['Y'][df_diet['X'] > 2e5]/10**func(np.log10(df_diet['X'][df_diet['X'] > 2e5])))**(1/3),
                'X_10_8': df_diet['X.1'][df_diet['X.1'] > 2e5], 'Y_10_8': df_diet['Y.1'][df_diet['X.1'] > 2e5], 'Yp_10_8': (df_diet['Y.1'][df_diet['X.1'] > 2e5]/10**func(np.log10(df_diet['X.1'][df_diet['X.1'] > 2e5])))**(1/3),
                'X_8_6': df_diet['X.2'][df_diet['X.2'] > 2e5], 'Y_8_6': df_diet['Y.2'][df_diet['X.2'] > 2e5], 'Yp_8_6': (df_diet['Y.2'][df_diet['X.2'] > 2e5]/10**func(np.log10(df_diet['X.2'][df_diet['X.2'] > 2e5])))**(1/3),
                'X_6_4': df_diet['X.3'][df_diet['X.3'] > 2e5], 'Y_6_4': df_diet['Y.3'][df_diet['X.3'] > 2e5], 'Yp_6_4': (df_diet['Y.3'][df_diet['X.3'] > 2e5]/10**func(np.log10(df_diet['X.3'][df_diet['X.3'] > 2e5])))**(1/3),
                'X_4_2': df_diet['X.4'][df_diet['X.4'] > 2e5], 'Y_4_2': df_diet['Y.4'][df_diet['X.4'] > 2e5], 'Yp_4_2': (df_diet['Y.4'][df_diet['X.4'] > 2e5]/10**func(np.log10(df_diet['X.4'][df_diet['X.4'] > 2e5])))**(1/3),
                'X_2_0': df_diet['X.5'][df_diet['X.5'] > 2e5], 'Y_2_0': df_diet['Y.5'][df_diet['X.5'] > 2e5], 'Yp_2_0': (df_diet['Y.5'][df_diet['X.5'] > 2e5]/10**func(np.log10(df_diet['X.5'][df_diet['X.5'] > 2e5])))**(1/3),})

                
label = {'gb': 'Spheres', 'ng': 'Natural gravel (NG) 1', 'gg': 'Rounded chips', 'ls': 'Rectangular prisms', 'oc': 'Faceted ellipsoids',
        'dn': 'Tempered glass', 'jo': 'NG 2', 'mb': 'NG 3', 'mk': 'NG 4', 'pg': 'NG 5', 'sf': 'Shell fragments'}
color = {'gb': sns.xkcd_rgb['cobalt blue'], 'gg': sns.xkcd_rgb['cerulean'], 'oc': sns.xkcd_rgb['blue green'], 'ng': sns.xkcd_rgb['maize'], 'ls': sns.xkcd_rgb['blood orange'], 'bm': 'y',
         'dn': sns.xkcd_rgb['cyan'], 'jo': sns.xkcd_rgb['lime green'], 'mb': sns.xkcd_rgb['black'], 'mk': sns.xkcd_rgb['pink'], 'pg': sns.xkcd_rgb['dark red'], 'sf': sns.xkcd_rgb['dusty purple']}
marker = {'gb': 'o', 'gg': 'D', 'oc': '*', 'ng': 'v', 'ls': 's',
        'dn': 'o', 'jo': 'D', 'mb': '*', 'mk': 'v', 'pg': 's', 'sf': 'o'}
size = {'gb': 4, 'gg': 4, 'oc': 8, 'ng': 6, 'ls': 4,
        'dn': 4, 'jo': 4, 'mb': 8, 'mk': 6, 'pg': 4, 'sf': 4}
grains = ['gb', 'oc', 'gg', 'ng', 'ls']
grains_2 = ['dn', 'jo', 'mb', 'mk', 'pg', 'sf']
grains_new = {'dn': 'tg', 'jo': 'ng2', 'mb': 'ng3', 'mk': 'ng4', 'pg': 'ng5', 'sf': 'sf'}

df_shape_2022 = pd.read_csv('1_grain_properties/extra_granular_materials_2022/2_var_shape_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_SV_2022 = pd.read_csv('1_grain_properties/extra_granular_materials_2022/4_var_settling_velocities_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_SV_VES_2022 = pd.read_csv('1_grain_properties/extra_granular_materials_2022/5_var_VES_settling_velocity_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_AoR_2022 = pd.read_csv('1_grain_properties/extra_granular_materials_2022/3_var_ARS_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv('1_grain_properties/2_var_shape.csv').set_index('Unnamed: 0').rename_axis(None)
df_SV = pd.read_csv('1_grain_properties/4_var_settling_velocities.csv').set_index('Unnamed: 0').rename_axis(None)
df_SV_VES = pd.read_csv('1_grain_properties/5_var_VES_settling_velocity.csv').set_index('Unnamed: 0').rename_axis(None)
df_AoR = pd.read_csv('1_grain_properties/3_var_ARS.csv').set_index('Unnamed: 0').rename_axis(None)


#%%
plot_xtra_grains_flag = True
fig, (ax2, ax1) = plt.subplots(1, 2, figsize=(18*cm, 6*cm))

# plot background data for axis 1
x_val = {'Yp_10': 1, 'Yp_10_8': .9, 'Yp_8_6': .7, 'Yp_6_4': .5, 'Yp_4_2': .3, 'Yp_2_0': .1}
fac = 0.03
for name in ['Yp_10', 'Yp_10_8', 'Yp_8_6', 'Yp_6_4', 'Yp_4_2', 'Yp_2_0']:
    y = xf[name].values
    if name is 'Yp_10':
        ax1.plot((x_val[name]*np.random.rand(y.size) * fac + x_val[name] - fac/2), 1/y**2, '.', c='k', label='Ref. 24', alpha=.3, markeredgewidth=0)
    else:
        ax1.plot((x_val[name]*np.random.rand(y.size) * fac + x_val[name] - fac/2), 1/y**2, '.', c='k', alpha=.3, markeredgewidth=0)

ax1.set_xlabel(r'Corey shape factor, $S_f = c/\sqrt{ab}$')
ax1.set_ylabel('Relative coeff. of drag,\n$C_{D_{settle}}/C_{o}$')
ax1.set_xscale('linear')
ax1.set_yscale('log')
ax1.set_xlim(0,1.05)
ax1.set_ylim(2e-1,10)
sns.despine()

# plot background data for axis 2
df_dai = pd.read_csv('other_shape_data/dai_robinson_data.csv', header=0)
x_dai = df_dai['circularity_dai'].values[1:].astype(np.float) * 0.3 + 0.7 # data adjusted to account for accidentally mislabelled axis when copying data from figure
ax2.plot(x_dai, np.tan(np.radians(df_dai['Unnamed: 3'].values[1:].astype(np.float))), 'D', markersize=3, c='k', label='Ref. 21', alpha=0.3, markeredgewidth=0)

df_carrigy = pd.read_csv('other_shape_data/carrigy.csv', header=1)
circ = (df_carrigy.X/100) * np.pi/4 + (1-df_carrigy.X/100)
ang = df_carrigy.Y / 3 + 20 
ax2.plot(circ, np.tan(np.radians(ang)), 'o', markersize=3, c='k', label='Ref. 22', alpha=0.3, markeredgewidth=0)

df_robinson = pd.read_csv('other_shape_data/robinson_circularity.csv', header=0)
ax2.plot(df_robinson['grains'].values[1:].astype(np.float), np.tan(np.radians(df_robinson['Unnamed: 1'].values[1:].astype(np.float))), '*', c='k', label='Ref. 23', alpha=0.3, markeredgewidth=0)
ax2.plot(df_robinson['sand_sphere_mix'].values[1:].astype(np.float), np.tan(np.radians(df_robinson['Unnamed: 3'].values[1:].astype(np.float))), '*', c='k', alpha=0.3, markeredgewidth=0)

fs = 7
fm = 0
mew = 0.25
ax2.legend(ncol=1, fontsize=fs, loc=3)
ax2.set_xlim(0.7,1.01)
ax2.set_ylim(0.2,1)
ax2.set_xscale('linear')
ax2.set_yscale('linear')
ax2.set_xlabel('Circularity, $S_c = 4\pi A/P^2$')
ax2.set_ylabel('Coeff. of static friction, $\mu_s$')
if plot_xtra_grains_flag is True:
    ax1.text(.95, 7, 'b', fontsize=fs+1, weight='bold')
    ax2.text(.99, .95, 'a', fontsize=fs+1, weight='bold')
else:
    ax1.text(1, 9, 'd', fontsize=fs+1, weight='bold')
    ax2.text(1.0, .95, 'c', fontsize=fs+1, weight='bold')
sns.despine()

# plot new grain data (2022)
if plot_xtra_grains_flag is True:
    for grain_in in grains_2:
        grain = grains_new[grain_in]
        yerr = 2*((df_SV_VES_2022.ws_sphere[grain]/df_SV_2022.vel_mean[grain])**2)*np.sqrt((df_SV_VES_2022.dws_sphere[grain]/df_SV_VES_2022.ws_sphere[grain])**2 + (df_SV_2022.dvel_mean[grain]/df_SV_2022.vel_mean[grain])**2)/np.sqrt(df_SV.vels.size)
        ax1.errorbar(df_shape_2022.CSF[grain], (df_SV_VES_2022.ws_sphere[grain]/df_SV_2022.vel_mean[grain])**2, c='w', fmt=marker[grain_in], markersize=size[grain_in]+fm)
        ax1.errorbar(df_shape_2022.CSF[grain], (df_SV_VES_2022.ws_sphere[grain]/df_SV_2022.vel_mean[grain])**2, xerr=df_shape_2022.dCSF[grain], yerr=yerr, ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, color=color[grain_in], markeredgecolor='w', markeredgewidth=mew,  fmt=marker[grain_in], markersize=size[grain_in], label=label[grain_in])
        
    def ellipse_perm(a,b): return 4*a*ss.ellipe(1.0 - b**2/a**2)  
    def ellipse_area(a,b): return np.pi * a * b
    def ellipse_circ(a,b): return 4*np.pi*ellipse_area(a,b)/ellipse_perm(a,b)**2

    Circ2 = {grains_new[grain_in]: ellipse_circ(df_shape_2022.C[grains_new[grain_in]], df_shape_2022.A[grains_new[grain_in]]) for grain_in in grains_2} # Corey shape factor for grain materials with variable shapes
    for grain_in in grains_2:
        grain = grains_new[grain_in]
        dAoR = 1
        dCirc2 = 0.02
        ax2.errorbar(Circ2[grain], np.tan(np.radians(df_AoR_2022.mean_angle[grain])), c='w', fmt=marker[grain_in], markersize=size[grain_in]*0.8+0.5)
        ax2.errorbar(Circ2[grain], np.tan(np.radians(df_AoR_2022.mean_angle[grain])), xerr=dCirc2, yerr=np.tan(np.radians(df_AoR_2022.std_angle[grain])), ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, c=color[grain_in], markeredgecolor='w', markeredgewidth=mew, fmt=marker[grain_in], markersize=size[grain_in]*0.8)


Circ = {
        'ng': 0.816,
        'ls': 0.823,
        'gg': 0.812,
        'oc': np.mean([0.82, 0.98]),
        'gb': 1.
        }
dCirc = {
        'ng': 0.016,
        'ls': 0.0093,
        'gg': 0.021,
        'oc': np.std([0.2, 0.3]),
        'gb': 0
        }


# plot flume shape data
for grain in grains:
    ax2.errorbar(Circ[grain], np.tan(np.radians(df_AoR.mean_angle[grain])), c='w', fmt=marker[grain], markersize=size[grain]+fm)
    ax2.errorbar(Circ[grain], np.tan(np.radians(df_AoR.mean_angle[grain])), xerr=dCirc[grain], yerr=np.tan(np.radians(df_AoR.std_angle[grain])), ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, c=color[grain], markeredgecolor='w', markeredgewidth=mew, fmt=marker[grain], markersize=size[grain])
    
for grain in grains:
    yerr = 2*((df_SV_VES.ws_sphere[grain]/df_SV.vel_mean[grain])**2)*np.sqrt((df_SV_VES.dws_sphere[grain]/df_SV_VES.ws_sphere[grain])**2 + (df_SV.dvel_mean[grain]/df_SV.vel_mean[grain])**2)
    ax1.errorbar(df_shape.CSF[grain], (df_SV_VES.ws_sphere[grain]/df_SV.vel_mean[grain])**2, c='w', fmt=marker[grain], markersize=size[grain]+fm)
    ax1.errorbar(df_shape.CSF[grain], (df_SV_VES.ws_sphere[grain]/df_SV.vel_mean[grain])**2, xerr=df_shape.dCSF[grain], yerr=yerr, ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, color=color[grain], markeredgecolor='w', markeredgewidth=mew,  fmt=marker[grain], markersize=size[grain], label=label[grain])

if plot_xtra_grains_flag is True: ax1.legend(ncol=3, fontsize=fs-2, loc=3)
else: ax1.legend(ncol=1, fontsize=fs-2, loc=3)

for grain in grains:
    ax2.errorbar(Circ[grain], np.tan(np.radians(df_AoR.mean_angle[grain])), c='w', fmt=marker[grain], markersize=size[grain]+fm)
    ax2.errorbar(Circ[grain], np.tan(np.radians(df_AoR.mean_angle[grain])), xerr=dCirc[grain], yerr=np.tan(np.radians(df_AoR.std_angle[grain])), ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, c=color[grain], markeredgecolor='w', markeredgewidth=mew, fmt=marker[grain], markersize=size[grain])


plt.tight_layout()
if plot_xtra_grains_flag is False: plt.savefig('Figure_1.jpg', dpi=300)
if plot_xtra_grains_flag is True: plt.savefig('Deal_EDfig4.jpg', dpi=300)

# %%
fig, (ax2, ax1, ax3) = plt.subplots(1, 3, figsize=(18*cm, 6*cm))
aor_gb, bed_angle = 24., 3.5
fm = 0.25
# plot new grain data (2022)
for grain_in in grains_2:
    grain = grains_new[grain_in]
    cstar = ((df_SV_VES_2022.ws_sphere[grain])/(df_SV_2022.vel_mean[grain]))**2 * df_shape_2022.CSF[grain]
    dcstar = cstar * np.sqrt((2*df_SV_2022.dvel_mean[grain]/df_SV_2022.vel_mean[grain])**2 + (df_shape_2022.dCSF[grain]/df_shape_2022.CSF[grain])**2)
    mugb = (np.tan(np.radians(aor_gb)) - np.tan(np.radians(bed_angle)))
    mustar = (np.tan(np.radians(df_AoR_2022.mean_angle[grain])) - np.tan(np.radians(bed_angle)))/mugb
    dmustar = ((1 + np.tan(np.radians(df_AoR_2022.mean_angle[grain]))**2)/mugb) * np.radians(df_AoR_2022.mean_angle[grain]) * (df_AoR_2022.std_angle[grain]/df_AoR_2022.mean_angle[grain])
    n = df_shape_2022.n[grain]
    ratio = cstar/mustar
    dratio = ratio * np.sqrt((dcstar/cstar)**2 + (dmustar/mustar)**2)

    ax1.errorbar(df_shape_2022.CSF[grain], cstar, c='w', fmt=marker[grain_in], markersize=size[grain_in]+fm)
    ax1.errorbar(df_shape_2022.CSF[grain], cstar, xerr=df_shape_2022.dCSF[grain], yerr=dcstar, ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, color=color[grain_in], markeredgecolor='w', markeredgewidth=mew,  fmt=marker[grain_in], markersize=size[grain_in], label=label[grain_in])

    ax2.errorbar(df_shape_2022.CSF[grain], mustar, c='w', fmt=marker[grain_in], markersize=size[grain_in]+fm)
    ax2.errorbar(df_shape_2022.CSF[grain], mustar, xerr=df_shape_2022.dCSF[grain], yerr=dmustar, ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, c=color[grain_in], markeredgecolor='w', markeredgewidth=mew, fmt=marker[grain_in], markersize=size[grain_in])

    ax3.errorbar(df_shape_2022.CSF[grain], ratio, c='w', fmt=marker[grain_in], markersize=size[grain_in]+fm)
    ax3.errorbar(df_shape_2022.CSF[grain], ratio, xerr=df_shape_2022.dCSF[grain], yerr=dratio, ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, color=color[grain_in], markeredgecolor='w', markeredgewidth=mew,  fmt=marker[grain_in], markersize=size[grain_in])

    tval = (ratio-1)/(dratio)
    dof = n-1
    pval = 2*(stats.t.cdf(-abs(tval), dof))
    
for grain in grains:
    cstar = ((df_SV_VES.ws_sphere[grain])/(df_SV.vel_mean[grain]))**2 * df_shape.CSF[grain]
    dcstar = cstar * np.sqrt((2*df_SV.dvel_mean[grain]/df_SV.vel_mean[grain])**2 + (df_shape.dCSF[grain]/df_shape.CSF[grain])**2)
    mugb = (np.tan(np.radians(aor_gb)) - np.tan(np.radians(bed_angle)))
    mustar = (np.tan(np.radians(df_AoR.mean_angle[grain])) - np.tan(np.radians(bed_angle)))/mugb
    dmustar = ((1 + np.tan(np.radians(df_AoR.mean_angle[grain]))**2)/mugb) * np.radians(df_AoR.mean_angle[grain]) * (df_AoR.std_angle[grain]/df_AoR.mean_angle[grain])
    n = df_shape.n[grain]
    ratio = cstar/mustar
    dratio = ratio * np.sqrt((dcstar/cstar)**2 + (dmustar/mustar)**2)

    ax1.errorbar(df_shape.CSF[grain], cstar, c='w', fmt=marker[grain], markersize=size[grain]+fm)
    ax1.errorbar(df_shape.CSF[grain], cstar, xerr=df_shape.dCSF[grain], yerr=dcstar, ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, color=color[grain], markeredgecolor='w', markeredgewidth=mew,  fmt=marker[grain], markersize=size[grain])

    ax2.errorbar(df_shape.CSF[grain], mustar, c='w', fmt=marker[grain], markersize=size[grain]+fm)
    ax2.errorbar(df_shape.CSF[grain], mustar, xerr=df_shape.dCSF[grain], yerr=dmustar, ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, c=color[grain], markeredgecolor='w', markeredgewidth=mew, fmt=marker[grain], markersize=size[grain], label=label[grain])

    ax3.errorbar(df_shape.CSF[grain], ratio, c='w', fmt=marker[grain], markersize=size[grain]+fm)
    ax3.errorbar(df_shape.CSF[grain], ratio, xerr=df_shape.dCSF[grain], yerr=dratio, ecolor=sns.xkcd_rgb['grey'], elinewidth=.9, capthick=.9, capsize=1, color=color[grain], markeredgecolor='w', markeredgewidth=mew,  fmt=marker[grain], markersize=size[grain])
    tval = (ratio-1)/(dratio)
    dof = n-1
    pval = 2*(stats.t.cdf(-abs(tval), dof))

x = np.linspace(0, 1)
ax3.plot(x, np.ones_like(x), 'k--', lw=0.8)

ax2.legend(ncol=1, fontsize=fs-2, loc=3)
ax1.legend(ncol=2, fontsize=fs-2, loc=3)

ax1.set_xlabel(r'Corey shape factor, $S_f$')
ax1.set_ylabel('Normalized coeff. of drag, $C^*$')
ax2.set_xlabel(r'Corey shape factor, $S_f$')
ax2.set_ylabel('Normalized coeff. of friction, $\mu^*$')
ax3.set_xlabel(r'Corey shape factor, $S_f$')
ax3.set_ylabel('$C^*/\mu^*$')

ax1.set_xlim(0,1.05)
ax1.set_ylim(0.5,2.6)
ax2.set_xlim(0,1.05)
ax2.set_ylim(0.5,2.6)
ax3.set_xlim(0,1.05)
ax3.set_ylim(0.25,1.75)
ax2.text(.1, 2.4,'g', fontsize=fs+1, weight='bold')
ax1.text(.1, 2.4,'h', fontsize=fs+1, weight='bold')
ax3.text(.1, 1.6,'i', fontsize=fs+1, weight='bold')

sns.despine()
plt.tight_layout()
plt.savefig('Deal_EDfig3_g-i.jpg', dpi=300)
# %%
fig, ax1 = plt.subplots(1,1,figsize=(12*cm, 9*cm))

yerr_cd = {grain: 2/((df_SV['vel_mean'][grain]/df_SV_VES.ws_sphere[grain]))**3 * ((df_SV['vel_mean'][grain]/df_SV_VES.ws_sphere[grain]) * np.sqrt((df_SV['dvel_mean'][grain]/df_SV['vel_mean'][grain])**2 + (df_SV_VES['dws_sphere'][grain]/df_SV_VES['ws_sphere'][grain])**2))/np.sqrt(df_SV.vels.size) for grain in grains}

df_shape_2022 = pd.read_csv('1_grain_properties/extra_granular_materials_2022/2_var_shape_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_SV_2022 = pd.read_csv('1_grain_properties/extra_granular_materials_2022/4_var_settling_velocities_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_SV_VES_2022 = pd.read_csv('1_grain_properties/extra_granular_materials_2022/5_var_VES_settling_velocity_2022.csv').set_index('Unnamed: 0').rename_axis(None)

for grain in grains: ax1.errorbar(df_shape.CSF[grain], (df_SV['vel_mean'][grain]/df_SV_VES.ws_sphere[grain])**(-2), yerr=yerr_cd[grain], c='w', fmt=marker[grain], markersize=size[grain])
for grain in grains: ax1.errorbar(df_shape.CSF[grain], (df_SV['vel_mean'][grain]/df_SV_VES.ws_sphere[grain])**(-2), yerr=yerr_cd[grain], xerr=df_shape.dCSF[grain], c=color[grain], fmt=marker[grain], ecolor=sns.xkcd_rgb['grey'], elinewidth=.5, capthick=.9, capsize=1, markeredgecolor='w', markeredgewidth=mew, markersize=size[grain], label=label[grain])

for grain_in in grains_2: 
    grain = grains_new[grain_in]
    ax1.errorbar(df_shape_2022.CSF[grain], (df_SV_2022.vel_mean[grain]/df_SV_VES_2022.ws_sphere[grain])**(-2), yerr=2*(df_SV_2022.vel_mean[grain]/df_SV_VES_2022.ws_sphere[grain])**(-2)*np.sqrt(((df_SV_2022.dvel_mean[grain]/df_SV_2022.vel_mean[grain])**(2)+(df_SV_VES_2022.dws_sphere[grain]/df_SV_VES_2022.ws_sphere[grain])**(2))), c='w', fmt=marker[grain_in], markersize=size[grain_in])
    ax1.errorbar(df_shape_2022.CSF[grain], (df_SV_2022.vel_mean[grain]/df_SV_VES_2022.ws_sphere[grain])**(-2), yerr=2*(df_SV_2022.vel_mean[grain]/df_SV_VES_2022.ws_sphere[grain])**(-2)*np.sqrt(((df_SV_2022.dvel_mean[grain]/df_SV_2022.vel_mean[grain])**(2)+(df_SV_VES_2022.dws_sphere[grain]/df_SV_VES_2022.ws_sphere[grain])**(2))), xerr=df_shape_2022.dCSF[grain], c=color[grain_in], fmt=marker[grain_in], ecolor=sns.xkcd_rgb['grey'], elinewidth=.5, capthick=.9, capsize=1, markeredgecolor='w', markeredgewidth=mew, markersize=size[grain_in], label=label[grain_in])


df = pd.read_csv('other_shape_data/ws_wn_csf.csv')
sf = np.linspace(0.2,1,100)
ax1.loglog(df['p=6'][1:].values.astype('float'), df['Unnamed: 1'][1:].values.astype('float')**(-2), '-', c='gray', label='Well rounded     ', lw=1)
# ax1.fill_between(df['p=6'][1:].values.astype('float'), np.ones(df['p=6'][1:].values.astype('float').size), df['Unnamed: 1'][1:].values.astype('float')**(-2), color='gray', alpha=.3)
ax1.plot(df['p=3.5'][1:].values.astype('float'), df['Unnamed: 3'][1:].values.astype('float')**(-2), '--', c='gray', label='Sub-angular    ', lw=1)
ax1.plot(df['p=2'][1:].values.astype('float'), df['Unnamed: 5'][1:].values.astype('float')**(-2), ':', c='gray', label='Very angular     ', lw=1)
ax1.plot((0,2),(1,1),'k--')

ax1.plot(sf, sf**-1, '--', c='r', label='$1/S_f$', lw=1)

sns.despine()
ax1.set_xlim(0.25, 1.05)
ax1.set_ylim(0.2, 10.1)
ax1.set_xlabel('Corey shape factor, $S_f$')
ax1.set_ylabel(r'$C_{D_{settle}}/C_{o}}$')
ax1.legend(ncol=3, loc=4, fontsize=fs)
ax1.text(.305, 1.3, 'Change in drag due to gross shape', fontsize=fs)
ax1.arrow(.33, 1.2, 0, -.1, head_width=0.01, head_length=0.1, fc='k', ec='k')
ax1.arrow(.43, 1.2, 0, -.1, head_width=0.0125, head_length=0.1, fc='k', ec='k')
ax1.arrow(.555, 1.2, 0, -.1, head_width=0.015, head_length=0.1, fc='k', ec='k')
ax1.arrow(.33, 1.55, 0, 1.3, head_width=0.01, head_length=0.3, fc='k', ec='k')
ax1.arrow(.43, 1.55, 0, .6, head_width=0.0125, head_length=0.25, fc='k', ec='k')
ax1.arrow(.555, 1.55, 0, .095, head_width=0.015, head_length=0.18, fc='k', ec='k')

plt.tight_layout()
plt.savefig('Deal_EDfig2.pdf', dpi=300)
# %%
